#include <ATen/cuda/cub.cuh>
#include <gtest/gtest.h>

TEST(NumBits, CubTest) {
  using at::cuda::cub::get_num_bits;
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000000000UL), 1);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000000001UL), 1);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000000010UL), 2);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000000011UL), 2);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000000100UL), 3);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000000111UL), 3);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000001000UL), 4);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000001111UL), 4);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000010000UL), 5);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000011111UL), 5);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000100000UL), 6);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000000111111UL), 6);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000001000000UL), 7);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000001111111UL), 7);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000010000000UL), 8);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000011111111UL), 8);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000100000000UL), 9);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000000111111111UL), 9);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000001000000000UL), 10);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000001111111111UL), 10);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000010000000000UL), 11);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000011111111111UL), 11);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000100000000000UL), 12);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000000111111111111UL), 12);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000001000000000000UL), 13);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000001111111111111UL), 13);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000010000000000000UL), 14);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000011111111111111UL), 14);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000100000000000000UL), 15);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000000111111111111111UL), 15);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000001000000000000000UL), 16);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000001111111111111111UL), 16);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000010000000000000000UL), 17);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000011111111111111111UL), 17);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000100000000000000000UL), 18);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000000111111111111111111UL), 18);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000001000000000000000000UL), 19);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000001111111111111111111UL), 19);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000010000000000000000000UL), 20);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000011111111111111111111UL), 20);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000100000000000000000000UL), 21);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000000111111111111111111111UL), 21);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000001000000000000000000000UL), 22);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000001111111111111111111111UL), 22);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000010000000000000000000000UL), 23);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000011111111111111111111111UL), 23);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000100000000000000000000000UL), 24);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000000111111111111111111111111UL), 24);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000001000000000000000000000000UL), 25);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000001111111111111111111111111UL), 25);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000010000000000000000000000000UL), 26);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000011111111111111111111111111UL), 26);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000100000000000000000000000000UL), 27);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000000111111111111111111111111111UL), 27);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000001000000000000000000000000000UL), 28);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000001111111111111111111111111111UL), 28);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000010000000000000000000000000000UL), 29);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000011111111111111111111111111111UL), 29);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000100000000000000000000000000000UL), 30);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000000111111111111111111111111111111UL), 30);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000001000000000000000000000000000000UL), 31);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000001111111111111111111111111111111UL), 31);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000010000000000000000000000000000000UL), 32);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000011111111111111111111111111111111UL), 32);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000100000000000000000000000000000000UL), 33);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000000111111111111111111111111111111111UL), 33);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000001000000000000000000000000000000000UL), 34);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000001111111111111111111111111111111111UL), 34);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000010000000000000000000000000000000000UL), 35);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000011111111111111111111111111111111111UL), 35);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000100000000000000000000000000000000000UL), 36);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000000111111111111111111111111111111111111UL), 36);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000001000000000000000000000000000000000000UL), 37);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000001111111111111111111111111111111111111UL), 37);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000010000000000000000000000000000000000000UL), 38);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000011111111111111111111111111111111111111UL), 38);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000100000000000000000000000000000000000000UL), 39);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000000111111111111111111111111111111111111111UL), 39);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000001000000000000000000000000000000000000000UL), 40);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000001111111111111111111111111111111111111111UL), 40);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000010000000000000000000000000000000000000000UL), 41);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000011111111111111111111111111111111111111111UL), 41);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000100000000000000000000000000000000000000000UL), 42);
  ASSERT_EQ(get_num_bits(0b0000000000000000000000111111111111111111111111111111111111111111UL), 42);
  ASSERT_EQ(get_num_bits(0b0000000000000000000001000000000000000000000000000000000000000000UL), 43);
  ASSERT_EQ(get_num_bits(0b0000000000000000000001111111111111111111111111111111111111111111UL), 43);
  ASSERT_EQ(get_num_bits(0b0000000000000000000010000000000000000000000000000000000000000000UL), 44);
  ASSERT_EQ(get_num_bits(0b0000000000000000000011111111111111111111111111111111111111111111UL), 44);
  ASSERT_EQ(get_num_bits(0b0000000000000000000100000000000000000000000000000000000000000000UL), 45);
  ASSERT_EQ(get_num_bits(0b0000000000000000000111111111111111111111111111111111111111111111UL), 45);
  ASSERT_EQ(get_num_bits(0b0000000000000000001000000000000000000000000000000000000000000000UL), 46);
  ASSERT_EQ(get_num_bits(0b0000000000000000001111111111111111111111111111111111111111111111UL), 46);
  ASSERT_EQ(get_num_bits(0b0000000000000000010000000000000000000000000000000000000000000000UL), 47);
  ASSERT_EQ(get_num_bits(0b0000000000000000011111111111111111111111111111111111111111111111UL), 47);
  ASSERT_EQ(get_num_bits(0b0000000000000000100000000000000000000000000000000000000000000000UL), 48);
  ASSERT_EQ(get_num_bits(0b0000000000000000111111111111111111111111111111111111111111111111UL), 48);
  ASSERT_EQ(get_num_bits(0b0000000000000001000000000000000000000000000000000000000000000000UL), 49);
  ASSERT_EQ(get_num_bits(0b0000000000000001111111111111111111111111111111111111111111111111UL), 49);
  ASSERT_EQ(get_num_bits(0b0000000000000010000000000000000000000000000000000000000000000000UL), 50);
  ASSERT_EQ(get_num_bits(0b0000000000000011111111111111111111111111111111111111111111111111UL), 50);
  ASSERT_EQ(get_num_bits(0b0000000000000100000000000000000000000000000000000000000000000000UL), 51);
  ASSERT_EQ(get_num_bits(0b0000000000000111111111111111111111111111111111111111111111111111UL), 51);
  ASSERT_EQ(get_num_bits(0b0000000000001000000000000000000000000000000000000000000000000000UL), 52);
  ASSERT_EQ(get_num_bits(0b0000000000001111111111111111111111111111111111111111111111111111UL), 52);
  ASSERT_EQ(get_num_bits(0b0000000000010000000000000000000000000000000000000000000000000000UL), 53);
  ASSERT_EQ(get_num_bits(0b0000000000011111111111111111111111111111111111111111111111111111UL), 53);
  ASSERT_EQ(get_num_bits(0b0000000000100000000000000000000000000000000000000000000000000000UL), 54);
  ASSERT_EQ(get_num_bits(0b0000000000111111111111111111111111111111111111111111111111111111UL), 54);
  ASSERT_EQ(get_num_bits(0b0000000001000000000000000000000000000000000000000000000000000000UL), 55);
  ASSERT_EQ(get_num_bits(0b0000000001111111111111111111111111111111111111111111111111111111UL), 55);
  ASSERT_EQ(get_num_bits(0b0000000010000000000000000000000000000000000000000000000000000000UL), 56);
  ASSERT_EQ(get_num_bits(0b0000000011111111111111111111111111111111111111111111111111111111UL), 56);
  ASSERT_EQ(get_num_bits(0b0000000100000000000000000000000000000000000000000000000000000000UL), 57);
  ASSERT_EQ(get_num_bits(0b0000000111111111111111111111111111111111111111111111111111111111UL), 57);
  ASSERT_EQ(get_num_bits(0b0000001000000000000000000000000000000000000000000000000000000000UL), 58);
  ASSERT_EQ(get_num_bits(0b0000001111111111111111111111111111111111111111111111111111111111UL), 58);
  ASSERT_EQ(get_num_bits(0b0000010000000000000000000000000000000000000000000000000000000000UL), 59);
  ASSERT_EQ(get_num_bits(0b0000011111111111111111111111111111111111111111111111111111111111UL), 59);
  ASSERT_EQ(get_num_bits(0b0000100000000000000000000000000000000000000000000000000000000000UL), 60);
  ASSERT_EQ(get_num_bits(0b0000111111111111111111111111111111111111111111111111111111111111UL), 60);
  ASSERT_EQ(get_num_bits(0b0001000000000000000000000000000000000000000000000000000000000000UL), 61);
  ASSERT_EQ(get_num_bits(0b0001111111111111111111111111111111111111111111111111111111111111UL), 61);
  ASSERT_EQ(get_num_bits(0b0010000000000000000000000000000000000000000000000000000000000000UL), 62);
  ASSERT_EQ(get_num_bits(0b0011111111111111111111111111111111111111111111111111111111111111UL), 62);
  ASSERT_EQ(get_num_bits(0b0100000000000000000000000000000000000000000000000000000000000000UL), 63);
  ASSERT_EQ(get_num_bits(0b0111111111111111111111111111111111111111111111111111111111111111UL), 63);
  ASSERT_EQ(get_num_bits(0b1000000000000000000000000000000000000000000000000000000000000000UL), 64);
  ASSERT_EQ(get_num_bits(0b1111111111111111111111111111111111111111111111111111111111111111UL), 64);
}
